import os
from celery import Celery
from rdflib import Dataset
# 处理视频脚本
from yolo_slowfast import main

# 配置 Celery
celery = Celery('tasks',
                broker="redis://localhost:6379/0",
                backend="redis://localhost:6379/0"
                )

# 调用yolo_slowfast脚本处理视频文件
def video_processing_task(input_path, output_path, config_args):
    class Args:
        def __init__(self, **entries):
            self.task_id = entries.get("task_id")  # 新增task_id参数
            self.output_dir = entries.get("output_dir", "results")  # 新增输出目录
            self.__dict__.update(entries)

    config = Args(
        input=input_path,
        output=output_path,
        **config_args
    )

    main(config)


import subprocess

# mp4转h264
def convert_video(input_path, output_path):
    command = [
        'ffmpeg',
        '-y',  # 覆盖输出文件
        '-i', input_path,
        '-c:v', 'libx264',
        '-profile:v', 'baseline',
        '-pix_fmt', 'yuv420p',
        '-movflags', '+faststart',
        output_path
    ]
    subprocess.run(command, check=True)


@celery.task(
    bind=True,
    autoretry_for=(Exception,),
    max_retries=2,
    retry_backoff=30,
    time_limit=3600,  # 1小时超时
    soft_timeout=300   # 5分钟无进展视为失败
)
def async_video_processing(self, task_id, input_path, output_path, config_args):
    try:
        config_args["task_id"] = task_id  # 添加task_id到配置参数
        # 立即更新为STARTED状态
        self.update_state(state='STARTED')
        print(f"任务状态已更新为STARTED: {self.AsyncResult(task_id).state}")
        print(f"task_id==>{task_id}")

        print(f"当前任务状态: {self.AsyncResult(self.request.id).state}")
        print(f"self.request.id==>{self.request.id}")

        # 进度回调包装器
        def progress_wrapper(current, total):
            self.update_state(
                state='PROGRESS',
                meta={
                    'status': 'processing',
                    'progress': int((current / total) * 100),
                    'current': current,
                    'total': total
                }
            )

        # 执行视频处理并传入回调
        video_processing_task(
            input_path,
            output_path,
            {**config_args, 'progress_callback': progress_wrapper}
        )

        # 修改临时文件命名（保持.mp4扩展名）
        temp_output = output_path.replace(".mp4", "_temp.mp4")  # 例如：output.mp4 → output_temp.mp4
        # 转换为H.264格式
        convert_video(output_path, temp_output)
        # 替换原文件
        os.replace(temp_output, output_path)

        # 验证输出文件
        if not os.path.exists(output_path):
            raise FileNotFoundError(f"输出文件 {output_path} 未生成")
        if os.path.getsize(output_path) == 0:
            raise ValueError(f"输出文件 {output_path} 为空")


        # 最终返回结果
        result = {
            'status': 'completed',
            'result_path': output_path,
            'progress': 100
        }

        self.update_state(state='SUCCESS', meta=result)  # 关键：通过meta传递完整结果
        print(f"最终任务状态: {self.AsyncResult(task_id).state}")
        return result

    except Exception as e:
        # 清理残留文件
        if os.path.exists(output_path):
            os.remove(output_path)
        self.update_state(state='FAILURE',
                          meta={'status': 'failed', 'error': str(e)})
        return {
            'status': 'failed',
            'error': str(e)
        }

# 示例YOLO训练脚本
"""
开始 模型培训 使用 train.py 脚本。基本论据包括

    --img:定义输入 图像大小 例如 --img 640).尺寸越大，精度越高，但需要的GPU 内存也越多。
    --batch:确定 批量大小 例如 --batch 16).选择GPU 可以处理的最大尺寸。
    --epochs:指定训练的总次数 纪元 例如 --epochs 100).一个历元代表对整个训练数据集的一次完整检测。
    --data:通往您的 dataset.yaml 文件（例如 --data coco128.yaml).
    --weights:初始权重文件的路径。使用预训练的权重（例如 --weights yolov5s.pt)，这样收敛速度更快，效果更好。要从头开始训练（不建议这样做，除非你有非常大的数据集和特殊需求），请使用 --weights '' --cfg yolov5s.yaml.

如果本地找不到预训练的权重，则会自动从最新的YOLOv5 版本中下载。

# Example: Train YOLOv5s on the COCO128 dataset for 3 epochs

遇到上述情况报错是由于模型batchsize设置的较大导致模型在训练时显卡内存不够。
遇到这种情况可以将batchsize减小。
在终端运行命令将batch修改合适的大小即可：
python train.py --img 320 --batch 4 --epochs 6

python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt
"""
"""
python train.py --data mask_data.yaml --cfg mask_yolov5s.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 4 --device cpu
"""
from yolov5 import train
def train_model(data_path, config, weights):
    train.run(
        weights=weights,  # 指定初始权重文件（随机初始化--weights ''）
        cfg=config.cfg,  # 选择model文件（models/yolo5s.yaml）
        data=config.data_path,  # 选择datase文件(data/coco128.yaml)
        epochs=config.epochs,  # 训练轮数
        batch_size=config.batch_size,  # 批量大小
        imgsz=config.imgsz,  # 输入图像大小
        lr0=config.lr0,  # 初始学习率
        # device='0',  # 使用GPU（如果有）
        # name='yolov5s_coco',  # 保存模型的名称
        # exist_ok=True, # 如果模型已存在，覆盖保存
    )

"""
    模型训练接口
"""
@celery.task(bind=True)
def train_yolo(self, dataset_id, config):
    dataset = Dataset.query.get(dataset_id)
    # 调用YOLO训练脚本
    command = [
        'python', 'train.py',
        '--data', os.path.join(dataset.path, 'data.yaml'),
        '--cfg', config['yolo_config'],
        '--weights', config['pretrained']
    ]
    subprocess.run(command, check=True)
    return {"status": "completed"}

"""
模型评估接口
"""
@celery.task(bind=True)
def evaluate_model(self, model_path, dataset_id):
    dataset = Dataset.query.get(dataset_id)
    # 调用评估脚本
    command = [
        'python', 'evaluate.py',
        '--weights', model_path,
        '--data', os.path.join(dataset.path, 'data.yaml')
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    # return {"accuracy": parse_accuracy(result.stdout)}
    return {"accuracy": 0.95}

# 测试任务
@celery.task(bind=True)
def test_task(self):
    # 更新状态为STARTED
    self.update_state(state='STARTED')
    print(f"测试任务状态: {self.AsyncResult(self.request.id).state}")
    return {"status": "ok"}